#include "proc.h"

// 测试用户程序的二进制代码（从 test_stub.bin 提取）
// 这个程序测试用户桩函数：直接调用 getpid() 和 write() 用户桩
static unsigned char test_program_bin[] = {
  0x37, 0x21, 0x00, 0x00, 0x41, 0x11, 0xa9, 0x42, 0x23, 0x23, 0x51, 0x00,
  0xb7, 0x62, 0x73, 0x74, 0x9b, 0x82, 0x42, 0x57, 0x23, 0x21, 0x51, 0x00,
  0xef, 0x00, 0x80, 0x07, 0x05, 0x45, 0x93, 0x05, 0x21, 0x00, 0x15, 0x46,
  0xef, 0x00, 0x40, 0x03, 0xef, 0x00, 0x80, 0x06, 0x01, 0xa0, 0x00, 0x00,
  0x85, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80, 0x89, 0x48, 0x73, 0x00,
  0x00, 0x00, 0x82, 0x80, 0x8d, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80,
  0x91, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80, 0x95, 0x48, 0x73, 0x00,
  0x00, 0x00, 0x82, 0x80, 0xc1, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80,
  0xd5, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80, 0x99, 0x48, 0x73, 0x00,
  0x00, 0x00, 0x82, 0x80, 0x9d, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80,
  0xa1, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80, 0xa5, 0x48, 0x73, 0x00,
  0x00, 0x00, 0x82, 0x80, 0xa9, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80,
  0xad, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80, 0xb1, 0x48, 0x73, 0x00,
  0x00, 0x00, 0x82, 0x80, 0xb5, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80,
  0xb9, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80, 0xbd, 0x48, 0x73, 0x00,
  0x00, 0x00, 0x82, 0x80, 0xc5, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80,
  0xc9, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80, 0xcd, 0x48, 0x73, 0x00,
  0x00, 0x00, 0x82, 0x80, 0xd1, 0x48, 0x73, 0x00, 0x00, 0x00, 0x82, 0x80,
  0x00, 0x00
};
static unsigned int test_program_bin_len = 218;

// 全局变量
struct cpu cpus[NCPU];
struct proc proc[NPROC];
struct proc *initproc;

int nextpid = 1;
spinlock pid_lock;

// 外部声明
extern char trampoline[]; 
extern pagetable_t kernel_pagetable;

// 帮助函数：等待父进程的wakeup不会被丢失
spinlock wait_lock;

// 为每个进程分配一个内核栈页面
// 在内核页表中映射它们到高地址
void proc_mapstacks(pagetable_t kpgtbl) {
  struct proc *p;
  
  for(p = proc; p < &proc[NPROC]; p++) {
    char *pa = alloc_page();
    if(pa == 0)
      panic("kalloc");
    uint64 va = KSTACK((int) (p - proc));
    map_page(kpgtbl, va, (uint64)pa, PTE_R | PTE_W);
  }
}

// 初始化进程表
void procinit(void) {
  struct proc *p;
  
  initlock(&pid_lock, "nextpid");
  initlock(&wait_lock, "wait_lock");
  
  for(p = proc; p < &proc[NPROC]; p++) {
    initlock(&p->lock, "proc");
    p->state = UNUSED;
    p->kstack = KSTACK((int) (p - proc));
  }
  
  printf("procinit: initialized process table\n");
}

// 加载测试程序到用户进程
// 将测试程序的二进制代码复制到用户进程的代码页中
static void load_test_program(struct proc *p) {
  extern pte_t* walk_lookup(pagetable_t pt, uint64 va);
  
  // 获取代码页的物理地址（在地址0处）
  uint64 code_va = 0;
  pte_t *pte;
  if((pte = walk_lookup(p->pagetable, code_va)) == 0)
    panic("load_test_program: code page not mapped");
  
  if((*pte & PTE_V) == 0)
    panic("load_test_program: code page not valid");
  
  uint64 code_pa = PTE2PA(*pte);
  char *code_page = (char*)code_pa;
  
  // 将测试程序复制到代码页
  memmove(code_page, test_program_bin, test_program_bin_len);
  
  p->trapframe->epc = 0;
}

// Set up first user process.
void
userinit(void)
{
  struct proc *p;

  p = allocproc();
  initproc = p;
  
 // p->cwd = namei("/");

  // 加载测试程序
  load_test_program(p);

  p->state = RUNNABLE;

  release(&p->lock);
}

// 分配一个进程结构体
// 如果找到未使用的进程，初始化内核运行所需的状态
// 返回时持有p->lock
struct proc* allocproc(void) {
  struct proc *p;

  for(p = proc; p < &proc[NPROC]; p++) {
    acquire(&p->lock);
    if(p->state == UNUSED) {
      goto found;
    } else {
      release(&p->lock);
    }
  }
  return 0;

found:
  p->pid = nextpid++;
  p->state = USED;

  // 分配陷阱帧页面
  if((p->trapframe = (struct trapframe *)alloc_page()) == 0){
    freeproc(p);
    release(&p->lock);
    return 0;
  }

  // 创建用户页表
  p->pagetable = proc_pagetable(p);
  if(p->pagetable == 0){
    freeproc(p);
    release(&p->lock);
    return 0;
  }

  // 设置新上下文，从forkret开始执行
  // forkret会返回到用户空间
  memset(&p->context, 0, sizeof(p->context));
  p->context.ra = (uint64)forkret;  
  p->context.sp = p->kstack + PGSIZE;  // 设置内核栈顶
  
  char *code_page = alloc_page();
  if(code_page == 0){
    freeproc(p);
    release(&p->lock);
    return 0;
  }


  if(map_page(p->pagetable, 0, (uint64)code_page, PTE_R | PTE_X | PTE_U) < 0){
    free_page(code_page);
    freeproc(p);
    release(&p->lock);
    return 0;
  }

  // 分配并映射用户栈
  char *stack_page = alloc_page();
  if(stack_page == 0){
    uvmunmap(p->pagetable, 0, 1, 0);
    free_page(code_page);
    freeproc(p);
    release(&p->lock);
    return 0;
  }
  memset(stack_page, 0, PGSIZE);

  if(map_page(p->pagetable, PGSIZE, (uint64)stack_page, PTE_R | PTE_W | PTE_U) < 0){
    uvmunmap(p->pagetable, 0, 1, 0);
    free_page(code_page);
    free_page(stack_page);
    freeproc(p);
    release(&p->lock);
    return 0;
  }

  p->trapframe->epc = 0;
  p->trapframe->sp = PGSIZE * 2;
  p->sz = PGSIZE * 2;
  
  return p;
}

// 释放进程结构体和相关数据
// 调用者必须持有p->lock
void freeproc(struct proc *p) {
  if(p->trapframe)
    free_page((void*)p->trapframe);
  p->trapframe = 0;
  if(p->pagetable)
    proc_freepagetable(p->pagetable, p->sz);
  p->pagetable = 0;
  p->sz = 0;
  p->pid = 0;
  p->parent = 0;
  p->name[0] = 0;
  p->chan = 0;
  p->killed = 0;
  p->xstate = 0;
  p->state = UNUSED;
}

// 为给定进程创建用户页表
// 映射蹦床和陷阱帧页面
pagetable_t proc_pagetable(struct proc *p) {
  pagetable_t pagetable;

  // 创建空页表
  pagetable = create_pagetable();
  if(pagetable == 0)
    return 0;

  // 映射陷阱帧（在页表顶部下方一页）
if(map_page(pagetable, TRAPFRAME, (uint64)(p->trapframe), PTE_R | PTE_W) < 0){
  destroy_pagetable(pagetable);
  return 0;
}

  // 映射蹦床代码（在页表顶部）
  if(map_page(pagetable, TRAMPOLINE, (uint64)trampoline, PTE_R | PTE_X) < 0){
    destroy_pagetable(pagetable);
    return 0;
  }




  return pagetable;
}

// 释放进程的页表
// 释放物理内存
void proc_freepagetable(pagetable_t pagetable, uint64 sz) {
  // 取消映射蹦床和陷阱帧（这些是共享的，不应该释放物理内存）
  // 只取消映射，不释放物理内存
  uvmunmap(pagetable, TRAMPOLINE, 1, 0);  // 取消映射蹦床
  uvmunmap(pagetable, TRAPFRAME, 1, 0);   // 取消映射陷阱帧
  
  // 释放用户地址空间的所有内存和页表结构
  uvmfree(pagetable, sz);
}

// forkret是fork进程第一次被scheduler调度时执行的地方
void forkret(void) {
  
  extern char userret[];
  static int first = 1;
  struct proc *p = myproc();

  // Still holding p->lock from scheduler.
  release(&p->lock);

  //  if (first) {
  //   // File system initialization must be run in the context of a
  //   // regular process (e.g., because it calls sleep), and thus cannot
  //   // be run from main().
  //   fsinit(ROOTDEV);

  //   first = 0;
  //   // ensure other cores see first=0.
  //   __sync_synchronize();

  //   // We can invoke kexec() now that file system is initialized.
  //   // Put the return value (argc) of kexec into a0.
  //   p->trapframe->a0 = kexec("/init", (char *[]){ "/init", 0 });
  //   if (p->trapframe->a0 == -1) {
  //     panic("exec");
  //   }
  // }

  // return to user space, mimicing usertrap()'s return.
  prepare_return();
  uint64 satp = MAKE_SATP(p->pagetable);
  uint64 trampoline_userret = TRAMPOLINE + (userret - trampoline);
  ((void (*)(uint64))trampoline_userret)(satp);
}



// 获取当前CPU的进程结构体指针
struct proc* myproc(void) {
  return cpus[0].proc;
}

// 复制虚拟内存页面（fork时使用）
// 返回复制的字节数，错误返回-1
int uvmcopy(pagetable_t old, pagetable_t new, uint64 sz) {
  pte_t *pte;
  uint64 pa, i;
  uint8 flags;
  
  // 遍历所有用户页面
  for(i = 0; i < sz; i += PGSIZE){
    if((pte = walk_lookup(old, i)) == 0)
      panic("uvmcopy: pte should exist");
    if((*pte & PTE_V) == 0)
      panic("uvmcopy: page not present");
    pa = PTE2PA(*pte);
    flags = PTE_FLAGS(*pte);
    
    // 分配新页面
    void *mem = alloc_page();
    if(mem == 0)
      goto err;
    
    // 复制页面内容
    memmove(mem, (char*)pa, PGSIZE);
    
    // 映射到新页表
    if(map_page(new, i, (uint64)mem, flags) != 0){
      free_page(mem);
      goto err;
    }
  }
  return 0;

 err:
  // 释放已分配的页面
  for(uint64 j = 0; j < i; j += PGSIZE){
    pte_t *pte_new = walk_lookup(new, j);
    if(pte_new && (*pte_new & PTE_V)){
      uint64 pa = PTE2PA(*pte_new);
      free_page((void*)pa);
    }
  }
  return -1;
}



// 将子进程交给init进程
// 调用者必须持有wait_lock
void reparent(struct proc *p) {
  struct proc *pp;
  
  for(pp = proc; pp < &proc[NPROC]; pp++){
    if(pp->parent == p){
      // 将子进程的父进程改为init
      pp->parent = initproc;
      // 唤醒init进程（可能需要处理这些子进程）
      wakeup(initproc);
    }
  }
}

// 唤醒所有在chan上等待的进程
void wakeup(void *chan) {
  struct proc *p;
  
  for(p = proc; p < &proc[NPROC]; p++){
    if(p != myproc()){
      acquire(&p->lock);
      if(p->state == SLEEPING && p->chan == chan) {
        p->state = RUNNABLE;
      }
      release(&p->lock);
    }
  }
}

// 在chan上睡眠，释放lk锁
void sleep(void *chan, spinlock *lk) {
  struct proc *p = myproc();
  
  // 必须持有p->lock才能改变state
  acquire(&p->lock);
  release(lk);

  // 进入睡眠
  p->chan = chan;
  p->state = SLEEPING;

  sched();

  // 醒来后，先重新获取lk锁，然后清理chan状态
  acquire(lk);
  p->chan = 0;
  
  // 释放p->lock
  release(&p->lock);
}

// 切换到调度器
void sched(void) {
  int intena;
  struct proc *p = myproc();

  if(!holding(&p->lock))
    panic("sched p->lock");
  if(p->state == RUNNING)
    panic("sched RUNNING");
  if(intr_get())
    panic("sched interruptible");

  intena = cpus[0].intena;
  swtch(&p->context, &cpus[0].context);
  cpus[0].intena = intena;
}

// 让出CPU：将当前进程状态改为RUNNABLE并切换到调度器
void yield(void) {
  struct proc *p = myproc();
  
  acquire(&p->lock);
  p->state = RUNNABLE;
  sched();
  release(&p->lock);
}

// 调度器：永不返回的循环
// 每个CPU调用这个函数来调度进程
// 实现轮转调度，确保公平性
void scheduler(void) {
  struct proc *p;
  struct cpu *c = &cpus[0];
 printf("scheduler: starting\n");
  c->proc = 0;
  c->last_proc_idx = 0;  // 初始化从0开始
  for(;;){
    // 避免死锁：如果有锁冲突，先开中断再关中断
    intr_on();
    intr_off();
    int found = 0;
    
    // 从上次调度的位置开始搜索（轮转调度）
    int start_idx = c->last_proc_idx;
    for(int i = 0; i < NPROC; i++) {
      int idx = (start_idx + i) % NPROC;
      p = &proc[idx];
      
      acquire(&p->lock);
      if(p->state == RUNNABLE) {
        //printf("scheduler: running process %d\n", p->pid);
        // 切换到选中的进程
        p->state = RUNNING;
        c->proc = p;
        c->last_proc_idx = (idx + 1) % NPROC;  // 更新下次起始位置
        swtch(&c->context, &p->context);
        // 进程运行完后切换回这里
        c->proc = 0;
        found = 1;
      }
      release(&p->lock);
      
      // 如果已经找到一个进程并执行完，跳出循环重新开始
      if(found)
        break;
    }
    
    if(found == 0) {
      // 没有可运行的进程，等待中断
      printf("scheduler: no runnable process found, waiting...\n");
      asm volatile("wfi");
    }
  }
}

// 创建新进程（fork）
int fork(void) {
  int i, pid;
  struct proc *np;
  struct proc *p = myproc();

  // 分配新进程
  if((np = allocproc()) == 0){
    return -1;
  }

  // 复制用户内存（如果sz > 0）
  if(p->sz > 0){
    if(uvmcopy(p->pagetable, np->pagetable, p->sz) < 0){
      freeproc(np);
      release(&np->lock);
      return -1;
    }
  }
  np->sz = p->sz;

  // 复制陷阱帧
  *(np->trapframe) = *(p->trapframe);

  // 子进程fork返回0
  np->trapframe->a0 = 0;

  // 复制文件描述符（暂不实现）
  for(i = 0; i < NOFILE; i++)
    if(p->ofile[i])
      np->ofile[i] = p->ofile[i];
  np->cwd = p->cwd;

  // 设置父进程
  pid = np->pid;
  
  acquire(&wait_lock);
  np->parent = p;
  release(&wait_lock);

  acquire(&np->lock);
  np->state = RUNNABLE;
  release(&np->lock);

  return pid;
}

// 退出当前进程
void exit(int status) {
  struct proc *p = myproc();
  if(p == 0)
    panic("exit");
  
  // 持有进程锁
  acquire(&p->lock);
  printf("exit: %d\n", p->pid);
  // 如果有init进程且当前进程是init，则panic
  if(initproc && p == initproc)
    panic("init exiting");

  // 关闭所有打开的文件（暂不实现文件系统，只清空指针）
  for(int fd = 0; fd < NOFILE; fd++){
    if(p->ofile[fd]){
      p->ofile[fd] = 0;
    }
  }

  // 开始退出：获取wait_lock
  acquire(&wait_lock);

  // 如果有init进程，将子进程交给init
  if(initproc) {
    reparent(p);
  }
  
  // 唤醒父进程（如果有的话）
  wakeup(p->parent);

  // 设置退出状态和进程状态
  p->xstate = status;
  p->state = ZOMBIE;

  // 释放wait_lock，但保留p->lock（sched()需要持有它）
  release(&wait_lock);
  
  // 切换到调度器（永不返回），scheduler会释放p->lock
  sched();
  panic("zombie exit");
}

// 等待子进程退出并回收
int wait(uint64 addr) {
  struct proc *pp;
  int havekids, pid;
  struct proc *p = myproc();

  acquire(&wait_lock);

  for(;;){
    // 扫描进程表查找ZOMBIE子进程
    havekids = 0;
    for(pp = proc; pp < &proc[NPROC]; pp++){
      if(pp->parent == p){
        // 确保子进程不在exit()或swtch()中
        acquire(&pp->lock);
        havekids = 1;
        if(pp->state == ZOMBIE){
          // 找到ZOMBIE子进程
          pid = pp->pid;
          if(addr != 0 && copyout(p->pagetable, addr, (char *)&pp->xstate,
                                  sizeof(pp->xstate)) < 0) {
            release(&pp->lock);
            release(&wait_lock);
            return -1;
          }
          // 释放ZOMBIE进程的资源
          freeproc(pp);
          release(&pp->lock);
          release(&wait_lock);
          return pid;
        }
        release(&pp->lock);
      }
    }

    // 如果没有子进程，返回-1
    if(!havekids){
      release(&wait_lock);
      return -1;
    }
    
    // 有子进程但不是ZOMBIE，等待子进程退出
    sleep(p, &wait_lock);
  }
}

// 复制数据到用户空间（简化版，暂时只做memcpy）
int copyout(pagetable_t pagetable, uint64 dstva, void *src, uint64 len) {
  // 简化实现：直接内存复制
  // 注意：在没有实现完整页表walk的情况下，这只是一个基本实现
  uint64 n;
  
  while(len > 0){
    n = len < PGSIZE ? len : PGSIZE;
    memmove((void*)dstva, src, n);
    len -= n;
    src += n;
    dstva += n;
  }
  return 0;
}

// 从用户空间复制数据到内核空间
int copyin(pagetable_t pagetable, char *dst, uint64 srcva, uint64 len) {
  uint64 va0, pa0;
  pte_t *pte;
  
  while(len > 0){
    va0 = PGROUNDDOWN(srcva);
    pte = walk_lookup(pagetable, va0);
    if(pte == 0 || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
      return -1;
    pa0 = PTE2PA(*pte);
    uint64 n = PGSIZE - (srcva - va0);
    if(n > len)
      n = len;
    memmove(dst, (void*)(pa0 + (srcva - va0)), n);
    len -= n;
    dst += n;
    srcva = va0 + PGSIZE;
  }
  return 0;
}

// Copy to either a user address, or kernel address,
// depending on usr_dst.
// Returns 0 on success, -1 on error.
int
either_copyout(int user_dst, uint64 dst, void *src, uint64 len)
{
  struct proc *p = myproc();
  if(user_dst){
    return copyout(p->pagetable, dst, src, len);
  } else {
    memmove((char *)dst, src, len);
    return 0;
  }
}

// Copy from either a user address, or kernel address,
// depending on usr_src.
// Returns 0 on success, -1 on error.
int
either_copyin(void *dst, int user_src, uint64 src, uint64 len)
{
  struct proc *p = myproc();
  if(user_src){
    return copyin(p->pagetable, dst, src, len);
  } else {
    memmove(dst, (char*)src, len);
    return 0;
  }
}

// 从用户空间复制字符串到内核空间
int copyinstr(pagetable_t pagetable, char *dst, uint64 srcva, uint64 max) {
  uint64 va0, pa0;
  pte_t *pte;
  int got_null = 0;
  
  while(got_null == 0 && max > 0){
    va0 = PGROUNDDOWN(srcva);
    pte = walk_lookup(pagetable, va0);
    if(pte == 0 || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
      return -1;
    pa0 = PTE2PA(*pte);
    uint64 n = PGSIZE - (srcva - va0);
    if(n > max)
      n = max;
    char *p = (char *)(pa0 + (srcva - va0));
    while(n > 0){
      if(*p == '\0'){
        *dst = '\0';
        got_null = 1;
        break;
      } else {
        *dst = *p;
      }
      --n;
      --max;
      p++;
      dst++;
    }
    srcva = va0 + PGSIZE;
  }
  if(got_null){
    return 0;
  } else {
    return -1;
  }
}



// 设置进程为killed状态
void setkilled(struct proc *p) {
  acquire(&p->lock);
  p->killed = 1;
  release(&p->lock);
}

// 检查进程是否被killed
int killed(struct proc *p) {
  int k;
  
  acquire(&p->lock);
  k = p->killed;
  release(&p->lock);
  return k;
}

// 创建一个执行指定任务的进程
// 这个方法允许你为不同的进程设置不同的任务入口点
struct proc* create_task_process(void (*task_func)(void)) {
  struct proc *p;
  
  // 分配进程结构体
  if((p = allocproc()) == 0) {
    return 0;
  }
  
  // 设置任务入口点（覆盖默认的forkret）
  p->context.ra = (uint64)task_func;
  
  // 释放allocproc时获取的锁
  release(&p->lock);
  
  // 设置为可运行状态
  acquire(&p->lock);
  p->state = RUNNABLE;
  release(&p->lock);
  return p;
}

void test_userProcess(void) {
  struct proc *p = allocproc();
  printf("userProcess: %d\n", p->pid);
  p->state = RUNNABLE;
  //fork();
  //exit(0);
}



// 简单的测试任务：只打印一些信息然后退出
void simple_task(void) {
  struct proc *p = myproc();
  
  release(&p->lock);
  
  printf("Process %d: Simple task started\n", p->pid);
  yield();
  exit(0);
}



